package async;

import org.apache.http.HttpResponse;
import org.apache.http.client.config.RequestConfig;
import org.apache.http.client.methods.HttpGet;
import org.apache.http.concurrent.FutureCallback;
import org.apache.http.impl.nio.client.CloseableHttpAsyncClient;
import org.apache.http.impl.nio.client.HttpAsyncClients;
import org.apache.http.impl.nio.conn.PoolingNHttpClientConnectionManager;
import org.apache.http.impl.nio.reactor.DefaultConnectingIOReactor;
import org.apache.http.nio.reactor.ConnectingIOReactor;
import org.apache.http.util.EntityUtils;
import org.apache.jmeter.config.Arguments;
import org.apache.jmeter.protocol.java.sampler.JavaSamplerClient;
import org.apache.jmeter.protocol.java.sampler.JavaSamplerContext;
import org.apache.jmeter.samplers.SampleResult;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.util.HashMap;
import java.util.Random;

public class AsyncSampler implements JavaSamplerClient {

	private final Logger log = LoggerFactory.getLogger(AsyncSampler.class);

	// 保存请求体
	public String url;
	public HashMap<String, String> headers = new HashMap<String, String>();
	public String body;

	// Jmeter界面参数设置：URL前缀/接口名/是否长连接
	public int MaxTotal;
	public int DefaultMaxPerRoute;
	public int SocketTimeout;
	public int ConnectTimeout;
	public int ConnectionRequestTimeout;

	public String MaxTotal_TEMPLATE_VALUE = "150";
	public String DefaultMaxPerRoute_TEMPLATE_VALUE = "150";
	public String SocketTimeout_TEMPLATE_VALUE = "250";
	public String ConnectTimeout_TEMPLATE_VALUE = "250";
	public String ConnectionRequestTimeout_TEMPLATE_VALUE = "1000";

	public ConnectingIOReactor ioReactor;
	public PoolingNHttpClientConnectionManager connectionManager;
	public CloseableHttpAsyncClient httpclient;
	public RequestConfig requestConfig;

	public AsyncRunInfo asyncRunInfo = AsyncRunInfo.getInstance();

	@Override
	public Arguments getDefaultParameters() {
		// TODO Auto-generated method stub

		Arguments params = new Arguments();

		params.addArgument("SetMaxTotal", MaxTotal_TEMPLATE_VALUE);
		params.addArgument("SetDefaultMaxPerRoute", DefaultMaxPerRoute_TEMPLATE_VALUE);
		params.addArgument("SetSocketTimeout", SocketTimeout_TEMPLATE_VALUE);
		params.addArgument("SetConnectTimeout", ConnectTimeout_TEMPLATE_VALUE);
		params.addArgument("SetConnectionRequestTimeout", ConnectionRequestTimeout_TEMPLATE_VALUE);

		return params;
	}

	@Override
	public void setupTest(JavaSamplerContext context) {
		// TODO Auto-generated method stub

		MaxTotal = context.getIntParameter("SetMaxTotal", 150);
		DefaultMaxPerRoute = context.getIntParameter("SetDefaultMaxPerRoute", 150);
		SocketTimeout = context.getIntParameter("SetSocketTimeout", 250);
		ConnectTimeout = context.getIntParameter("SetConnectTimeout", 250);
		ConnectionRequestTimeout = context.getIntParameter("SetConnectionRequestTimeout", 1000);

		log.info("AsyncSampler [MaxTotal=" + MaxTotal + ", DefaultMaxPerRoute=" + DefaultMaxPerRoute
				+ ", SocketTimeout=" + SocketTimeout + ", ConnectTimeout=" + ConnectTimeout
				+ ", ConnectionRequestTimeout=" + ConnectionRequestTimeout + "]");

		try {

			ioReactor = new DefaultConnectingIOReactor();
			connectionManager = new PoolingNHttpClientConnectionManager(ioReactor);
			connectionManager.setMaxTotal(MaxTotal);
			connectionManager.setDefaultMaxPerRoute(DefaultMaxPerRoute);

			httpclient = HttpAsyncClients.custom().setConnectionManager(connectionManager).build();
			httpclient.start();

			requestConfig = RequestConfig.custom().setSocketTimeout(SocketTimeout).setConnectTimeout(ConnectTimeout)
					.setConnectionRequestTimeout(ConnectionRequestTimeout).build();

			asyncRunInfo.start();

		} catch (Exception e) {
			// TODO: handle exception
			e.printStackTrace();
		}

	}

	@Override
	public SampleResult runTest(JavaSamplerContext context) {

		asyncRunInfo.sendCount();

		HttpGet request = new HttpGet("http://127.0.0.1:8080/zujizhuang?id=" + new Random().nextInt(100000000));
		request.setConfig(requestConfig);
		request.addHeader("connection", "keep-alive");

		SampleResult sampleResult = new SampleResult();
		sampleResult.setSamplerData("http://127.0.0.1:8080/zujizhuang");

		sampleResult.sampleStart();

		httpclient.execute(request, new FutureCallback<HttpResponse>() {

			public void completed(final HttpResponse response) {

				asyncRunInfo.receCount();

				try {
					String content = EntityUtils.toString(response.getEntity(), "UTF-8");
					if (!content.contains("HELLO")) {
						log.error(content);
					}

				} catch (IOException e) {
					e.printStackTrace();
				}

				request.releaseConnection();
			}

			public void failed(final Exception ex) {

				asyncRunInfo.failCount();
				// log.error("failed: " + ex.toString());

				request.releaseConnection();
			}

			public void cancelled() {

				asyncRunInfo.cancCount();
				// log.error("cancelled: ");

				request.releaseConnection();
			}
		});

		sampleResult.sampleEnd();

		sampleResult.setSuccessful(true);

		return sampleResult;
	}

	@Override
	public void teardownTest(JavaSamplerContext context) {

		try {
			httpclient.close();
		} catch (IOException e) {

			e.printStackTrace();
		}

		asyncRunInfo.stop();
	}
}
